import torch
import time
import torch_npu
from torch_npu.contrib import transfer_to_npu


def test(a, is_npu=False):
    start = time.time()
    for i in range(10):
        a = a @ a
    return time.time() - start


device = torch.device('npu')
# device = torch.device('cuda')


# 生成随机浮点数张量
print('随机生成张量')
random_tensor = torch.randn((1000, 1000), dtype=torch.float16)  # 使用标准正态分布生成随机张量

print('在CPU上测试')
print(f'时间：{test(random_tensor)}')

print('移动张量到npu')
# 将随机张量移动到NPU上
random_tensor = random_tensor.to(device)

print('在NPU上测试')
print(f'时间：{test(random_tensor)}')

# # 输出结果
# print("Result tensor shape:", result_tensor.shape)
# print("Result tensor example:", result_tensor[0][:5])  # 输出结果的前5个元素
